package com.senke.aiagent.rag;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.retrieval.join.ConcatenationDocumentJoiner;
import org.springframework.ai.rag.retrieval.join.DocumentJoiner;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 多查询扩展器
 */
@Component
@Slf4j
public class MyMultiQueryExpander {

    @Resource
    private ChatModel dashscopeChatModel;

    @Resource
    private VectorStore loveAppVectorStore;

    // 返回改写后的 prompt
    public List<Document> expandQuery(String message) {
        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel);
        MultiQueryExpander multiQueryExpander = MultiQueryExpander
                .builder()
                .chatClientBuilder(builder)
                .numberOfQueries(3)
                .includeOriginal(true)
                .build();
        Query query = new Query(message);
        List<Query> queryList = multiQueryExpander.expand(query);
        // 创建基于向量存储的文档检索器
        VectorStoreDocumentRetriever documentRetriever = VectorStoreDocumentRetriever
                .builder()
                .vectorStore(loveAppVectorStore)
                .similarityThreshold(0.5)
                .topK(5)
                .build();
        // 从多个查询或数据源获取的文档集合
        Map<Query, List<List<Document>>> documentsForQuery = new HashMap<>();
        queryList.forEach(q -> {
            log.info("QueryText: {}", q.text());
            List<Document> retrieve = documentRetriever.retrieve(q);
            documentsForQuery.put(q, List.of(retrieve));
        });
        // 创建文档合并器实例
        DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
        return documentJoiner.join(documentsForQuery);
    }
}
